from instruction_re.utils.evaluate_utils import (
    parse_simple_template_output_to_quintuples,
)

if __name__ == "__main__":
    # test_output = (
    #     "Hakawati Theatre (organization) <organization_based_in> Jerusalem (location)",
    #     "Hakawati Theatre (Organization)<organization_based_in> Jerusalem (Location)",
    # )
    test_output = (
        "University of Virginia (location) located in Charlottesville (location)"
    )

    label2token = {
        "organization based in": "<organization_based_in>",
        "located in": "located in",
    }

    print(
        parse_simple_template_output_to_quintuples(
            test_output,
            list(label2token.values()),
            ["person", "location", "organization"],
            sep="<sep>",
        )
    )

# import re


# def create_relation_regex(relation_types):
#     # 转义特殊字符并用 '|' 连接所有关系类型
#     escaped_relations = [re.escape(rel) for rel in relation_types]
#     relation_pattern = "|".join(escaped_relations)

#     # 构建完整的正则表达式
#     pattern = r"""
#     ^                           # 字符串开始
#     (\w+(?:\s\w+)*)             # 实体1: 一个或多个单词
#     \s+                         # 空白字符
#     \((\w+)\)                   # 实体1类型: 括号中的单词
#     \s+                         # 空白字符
#     <({})>                      # 关系类型: 尖括号中的任意一种关系
#     \s+                         # 空白字符
#     (\w+(?:\s\w+)*)             # 实体2: 一个或多个单词
#     \s+                         # 空白字符
#     \((\w+)\)                   # 实体2类型: 括号中的单词
#     $                           # 字符串结束
#     """.format(
#         relation_pattern
#     )

#     return re.compile(pattern, re.VERBOSE)


# # 示例使用
# relation_types = ["organization_based_in", "born_in", "employed_by"]
# regex = create_relation_regex(relation_types)

# # 测试正则表达式
# # test_string = (
# #     # "Havana (location) <organization_based_in> Radio Reloj Network (organization)"
# #     "Havana (location) <organization_based_in> Radio Reloj Network (organization)"
# # )
# test_string = (
#     "Havana (location) <organization_based_in> Radio Reloj Network (organization)"
# )
# match = regex.match(test_string)

# if match:
#     entity1, entity1_type, relation, entity2, entity2_type = match.groups()
#     print(f"Entity 1: {entity1} ({entity1_type})")
#     print(f"Relation: {relation}")
#     print(f"Entity 2: {entity2} ({entity2_type})")
# else:
#     print("No match found")
